import io
import json
from PIL import Image
import torch
from torchvision import transforms
from flask import Flask, jsonify, request
from resnet import resNet
import numpy as np
import cv2

# GPU CPU选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

net = resNet()
net.to(device)

# 加载模型
net.load_state_dict(torch.load(".\models\model_resnet"))

net.eval()

label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]

# 定义图像预处理函数
test_transform = transforms.Compose([
    transforms.CenterCrop((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 初始化Flask应用
app = Flask(__name__)


# 定义RESTful路由
@app.route('/predict', methods=['POST'])
def predict():
    # 获取上传的照片
    file = request.files['image']
    img_bytes = file.read()

    # 将照片转换为PIL图像
    img = Image.open(io.BytesIO(img_bytes))
    # 预处理图像
    img_tensor = test_transform(img)
    inputs = torch.unsqueeze(img_tensor, dim=0)
    inputs = inputs.to(device)
    outputs = net(inputs)
    _, prediction = torch.max(outputs.data, dim=1)
    # 返回预测结果
    result = {'class': str(label_name[prediction.cpu().numpy()[0]])}
    return jsonify(result)


# 运行Flask应用
if __name__ == '__main__':
    app.run(debug=True)
